//============================================================================
// Name         : sibefs.cpp
// Author       : Marcin Krystianc (marcin.krystianc@gmail.com)
// Version      : 1.2
// License      : GPL
// URL          : http://code.google.com/p/sibefs/
// Description  : Simple Filesystem Benchmark
//============================================================================

// compile:
// Linux32: g++ -Wall -m32 -O3 sibefs.cpp -o"sibefs32"
// Linux64: g++ -Wall -m64 -O3 sibefs.cpp -o"sibefs64"
// Windows: g++ -Wall -O3 sibefs.cpp -o"sibefs" -D SIBEFS_WIN

#include <iostream>
#include <string>
#include <fstream>
#include <ctime>
#include <sstream>
#include <iomanip>
#include <cstdlib>
#include <list>
#include <algorithm>

#include <signal.h>

#include <unistd.h>
#include <fcntl.h>
#include <sys/time.h>

using namespace std;
typedef unsigned long long int U_LONG;

const string fName1("a_sibefs.");
const string fName2("b_sibefs.");

U_LONG MAX_MEM      = U_LONG(1) << 27;  // max memory usage =128 MB
U_LONG MAX_FILE     = U_LONG(1) << 32;  // max file size to test
U_LONG MIN_FILE     = U_LONG(1) << 10;  // min file size to test
U_LONG HDD_LIMIT    = U_LONG(1) << 28;  // maximum disk space to take in tests
U_LONG REPEATS      = U_LONG(1);        // repeats
U_LONG FSYNC        = U_LONG(1);        // use fsync ?

string DIR;             // directory to test , default is executable file directory
unsigned char *GBUF;    // global buffer for files operation


#ifdef WIN32
    #define SIBEFS_WIN
#endif


#ifdef SIBEFS_WIN
    #include <sys\stat.h>
    #define OPENW_FLAGS (O_CREAT  | O_WRONLY | _O_BINARY)
    #define OPENR_FLAGS (O_RDONLY | _O_BINARY )
    #define OPENW_MODE  (_S_IREAD | _S_IWRITE)
#else
    #define OPENW_FLAGS (O_CREAT  | O_WRONLY  	| O_LARGEFILE)
    #define OPENR_FLAGS (O_RDONLY | O_LARGEFILE )
    #define OPENW_MODE  (S_IWUSR  | S_IRUSR)
#endif

//============================================================================
void initBuf(unsigned char * buf, U_LONG size);
void benchWriteFile(string fName, U_LONG fsize);
U_LONG benchReadFile(string fName);
U_LONG benchCopyFile(string fSName, string fDName);
bool isFile(string fName);
string fileSizeToString(U_LONG bytes);
string fileSpeedToString(U_LONG bytes, U_LONG msces);
U_LONG getTime();
bool parseCmd(int argc, char* argv[]);
U_LONG getCmdVal(const char* str);
bool printParameters();
int fsyncOsWrapper(int fd);

//============================================================================
int fsyncOsWrapper(int fd)
{
#ifdef SIBEFS_WIN
    return _commit (fd);
#else
    return fsync(fd);
#endif
}

//============================================================================
class cFileCloser
{
public:
    cFileCloser (int _fd):fd(_fd){};
    ~cFileCloser ()     {close (fd);};
protected:
    int fd;
};
//============================================================================
class cFileCleaner
{
public:
    ~cFileCleaner();
    void registerFile(string fName);
    void unregisterFile(string fName);
protected:
    list<string> fileList;
} fileCleanerA, fileCleanerB;

//============================================================================
cFileCleaner::~cFileCleaner()
{
    for (list<string>::iterator it = fileList.begin(); it != fileList.end(); ++it)
    {
        remove(it->c_str());
        cout << "~cFileCleaner():" << *it << endl;
    }
}

//============================================================================
void cFileCleaner::registerFile(string fName)
{
    fileList.push_back(fName);
}

//============================================================================
void cFileCleaner::unregisterFile(string fName)
{
    list<string>::iterator it = find(fileList.begin(), fileList.end(), fName);
    if (it != fileList.end())
    {
        fileList.erase(it);
        remove(fName.c_str());
    }
}

//============================================================================
void terminate(int param)
{
    cout << "Terminating program, some files may still remain on the disk !" << endl;
    exit(0); // will call destructor of fileCleanerA and fileCleanerB
}

//============================================================================
void initBuf(unsigned char * buf, U_LONG size)
{
    unsigned int *pBuf = (unsigned int *) buf;
    for (unsigned int i = size / sizeof(unsigned int); 0 < i; --i)
        *(pBuf++) += i;
}

//============================================================================
void benchWriteFile(string fName, U_LONG fsize)
{
    if (isFile(fName))
        throw string("benchWriteFile File exists: " + fName);

    int fdOut;
    if (-1 == (fdOut = open(fName.c_str(), OPENW_FLAGS, OPENW_MODE)))
        throw string("benchWriteFile Cannot open for write:" + fName);

    cFileCloser fdOutCloser (fdOut);

    for (U_LONG tmp = 0; fsize > 0; fsize -= tmp)
    {
        tmp = fsize > MAX_MEM ? MAX_MEM : fsize;

        if (-1 == write(fdOut, GBUF, tmp))
            throw string("benchWriteFile Write file error: " + fName);
    }

    if (FSYNC)
    if (-1 == fsyncOsWrapper(fdOut))
        throw string("benchWriteFile Write file error: " + fName);

}

//============================================================================
U_LONG benchReadFile(string fName)
{
    int fdIn;
    if (-1 == (fdIn = open(fName.c_str(), OPENR_FLAGS)))
        throw string("benchReadFile Cannot open for read: " + fName);

    cFileCloser fdInCloser (fdIn);

    U_LONG fSize = 0;
    int ret = 0;
    do
    {
        ret = read(fdIn, GBUF, MAX_MEM);
        if (ret == -1)
            throw string("benchReadFile Read file error: " + fName);

        fSize += ret;
    } while (ret);


    return fSize;
}

//============================================================================
U_LONG benchCopyFile(string fSName, string fDName)
{

    if (isFile(fDName))
        throw string("benchCopyFile File exists: " + fDName);

    int fdIn, fdOut;

    if (-1 == (fdIn = open(fSName.c_str(), OPENR_FLAGS)))
        throw string("benchCopyFile Cannot open for read: " + fSName);

    if (-1 == (fdOut = open(fDName.c_str(), OPENW_FLAGS, OPENW_MODE)))
        throw string("benchCopyFile Cannot open for write:" + fDName);

    cFileCloser fdInCloser(fdIn);
    cFileCloser fdOutClose(fdOut);


    U_LONG fSize = 0;
    int ret = 0;

    do
    {
        ret = read(fdIn, GBUF, MAX_MEM);
        if (ret == -1)
            throw string("benchCopyFile Read file error: " + fSName);

        if (-1 == write(fdOut, GBUF, ret))
            throw string("benchCopyFile Write file error: " + fDName);

        fSize += ret;
    } while (ret);

    if (FSYNC)
    if (-1 == fsyncOsWrapper(fdOut))
        throw string("benchCopyFile Write file error: " + fDName);

    return fSize;
}

//============================================================================
bool isFile(string fName)
{
    ifstream infile(fName.c_str(), ifstream::binary);
    return infile.good();
}

//============================================================================
string fileSizeToString(U_LONG bytes)
{
    stringstream ss;
    float fBytes = bytes;
    if (bytes < (1l << 10)) // less than KB
        ss << bytes << "  B";

    else if (bytes < (1l << 20)) // less than MB
        ss << fixed << setprecision(2) << fBytes / (1l << 10) << " KB";

    else if (bytes < (1l << 30)) // less than GB
        ss << fixed << setprecision(2) << fBytes / (1l << 20) << " MB";

    else
        ss << fixed << setprecision(2) << fBytes / (1l << 30) << " GB";

    return ss.str();
}

//============================================================================
string fileSpeedToString(U_LONG bytes, U_LONG msces)
{
    stringstream ss;

    if (msces == 0)
    {
        ss << setw(13) << "inf";
        return ss.str();
    }

    float fBytes = bytes;
    float speed = fBytes / msces * (1e3);

    if (speed < (1l << 10)) // less than KB
        ss << fixed << setprecision(2) << speed << "  B/s";

    else if (speed < (1l << 20)) // less than MB
        ss << fixed << setprecision(2) << speed / (1l << 10) << " KB/s";

    else if (speed < (1l << 30)) // less than GB
        ss << fixed << setprecision(2) << speed / (1l << 20) << " MB/s";

    else
        ss << fixed << setprecision(2) << speed / (1l << 30) << " GB/s";

    return ss.str();
}

//============================================================================
U_LONG getTime()
{
    timeval ts;
    gettimeofday(&ts,0);
    return (ts.tv_sec * 1000 + (ts.tv_usec / 1000));
}

//============================================================================
U_LONG getCmdVal(const char* str)
{
    char *pC;
    U_LONG tmp = strtoul(str, &pC, 10);

    switch (*pC)
    {
        case 'k':
        case 'K':
            return tmp * 1l << 10;

        case 'm':
        case 'M':
            return tmp * 1l << 20;

        case 'g':
        case 'G':
            return tmp * 1l << 30;
    }

    return tmp;
}

//============================================================================
bool parseCmd(int argc, char* argv[])
{
    int arg = 1;
    string s1, s2;

    while (arg < argc)
    {
        s1 = argv[arg++];
        if ("-i" == s1) // info
        {
            cout << "sizeof(ULONG):" << sizeof(U_LONG) << endl;
            exit(0);
        }

        if (arg == argc)
            return false;

        s2 = argv[arg++];

        if ("-d" == s1) // directory to test
            DIR = s2;

        if ("-m" == s1) // memory to use
            MAX_MEM = getCmdVal(s2.c_str());

        if ("-f" == s1) // minimum file size
            MIN_FILE = getCmdVal(s2.c_str());

        if ("-F" == s1) // max file size
            MAX_FILE = getCmdVal(s2.c_str());

        if ("-r" == s1) // repeat
            REPEATS = getCmdVal(s2.c_str());

        if ("-q" == s1) // quota - HDD limit
            HDD_LIMIT = getCmdVal(s2.c_str());

        if ("-s" == s1) // use fsync() ?
            FSYNC = getCmdVal(s2.c_str());
    }
    return true;
}

//============================================================================
bool printParameters()
{
    cout << "Sibefs";
    cout << " -m=" << fileSizeToString(MAX_MEM);
    cout << " -f=" << fileSizeToString(MIN_FILE);
    cout << " -F=" << fileSizeToString(MAX_FILE);
    cout << " -q=" << fileSizeToString(HDD_LIMIT);
    cout << " -r=" << REPEATS;
    cout << " -s=" << FSYNC;
    cout << " -d=" << DIR;

    cout << endl;

    if (!MAX_MEM || !MIN_FILE || !MAX_FILE || !HDD_LIMIT)
        return false;

    return true;
}

//============================================================================
int main(int argc, char* argv[])
{

    if (!parseCmd(argc, argv))
    {
        cout << "Sibefs - Simple Filesystem Benchmark" << endl;
        cout << "Options:" << endl;
        cout << "-d Directory to test, by default current directory is used." << endl;
        cout << "-m RAM to use by sibefs. This is size of one block for read/write command." << endl;
        cout << "-f Minimum file size. " << endl;
        cout << "-F Maximum file size. " << endl;
        cout << "-r Number of test repeats. " << endl;
        cout << "-q Disk space to use. " << endl;
        cout << "-s Use fsync() after write each file. " << endl;
        cout << endl;
        cout << "Benchmark scenario:" << endl;
        cout << "1. Create and write n=q/(size*2) files " << endl;
        cout << "2. Make a copy of n files " << endl;
        cout << "3. Read 2*n files " << endl;
        cout << "4. Delete 2*n files (original and copy)" << endl;
        cout << "5. Calculate size=0.5*size" << endl;
        cout << "6. Go to step 1" << endl;
        cout << endl;
        cout << "Examples:" << endl;
        cout << "sibefs -f 1k -F 10M -q 1G -m 1024 -s 0" << endl;
        cout << "sibefs -q 100m -s 1" << endl;
        cout << endl;
        cout << "No warranty. Use at Your own risk !" << endl;
        cout << endl;

        return 0;
    }

    if (!printParameters())
    {
        cout << "Invalid parameters !" << endl;
        return 0;
    }

    signal(SIGABRT, terminate);
    signal(SIGFPE, terminate);
    signal(SIGILL, terminate);
    signal(SIGINT, terminate);
    signal(SIGSEGV, terminate);
    signal(SIGTERM, terminate);

    GBUF = new unsigned char[MAX_MEM];
    if (!GBUF)
    {
        cout << "Cannot allocate " << MAX_MEM << " bytes" << endl;
        return 1;
    };

    try
    {
        initBuf(GBUF, MAX_MEM);

        while (REPEATS--)
        {
            U_LONG fCount = 0;
            U_LONG fSize = min(HDD_LIMIT / 2, MAX_FILE);

            for (; fSize >= MIN_FILE; fSize /= 2)
            {
                fCount = HDD_LIMIT / fSize / 2;

                U_LONG writeSize = 0, copySize = 0, readSize = 0;
                U_LONG writeTime = 0, copyTime = 0, readTime = 0, deleteTime =
                        0, completeTime = 0 , completeSeconds=0;

                cout << "Testing: " << setw(6) << fCount << "x" << setw(10)<< right << fileSizeToString(fSize);
                cout.flush();

                completeTime    = getTime();
                completeSeconds = time(0);

                // Perform write test
                cout << " | write:";
                cout.flush();
                writeTime = getTime();
                for (U_LONG i = 0; i < fCount; ++i)
                {
                    stringstream ss1, ss2;
                    ss1 << DIR << fName1 << i;
                    fileCleanerA.registerFile(ss1.str());
                    benchWriteFile(ss1.str(), fSize);
                    writeSize += fSize;
                }
                writeTime = getTime() - writeTime;
                cout << setw(12) << fileSizeToString(writeSize) << "," << setw(12) << right << fileSpeedToString(writeSize, writeTime);
                cout.flush();

                // Perform copy test
                cout << " | copy:";
                cout.flush();
                copyTime = getTime();
                for (U_LONG i = 0; i < fCount; ++i)
                {
                    stringstream ss1, ss2;
                    ss1 << DIR << fName1 << i;
                    ss2 << DIR << fName2 << i;
                    fileCleanerB.registerFile(ss2.str());
                    copySize += benchCopyFile(ss1.str(), ss2.str());
                }
                copyTime = getTime() - copyTime;
                cout << setw(12) << fileSizeToString(copySize) << "," << setw(12) << right << fileSpeedToString(copySize, copyTime);
                cout.flush();

                // Perform read test
                cout << " | read:";
                cout.flush();
                readTime = getTime();

                for (U_LONG i = 0; i < fCount; ++i)
                {
                    stringstream ss1, ss2;
                    ss1 << DIR << fName1 << i;
                    ss2 << DIR << fName2 << i;
                    readSize += benchReadFile(ss1.str());
                    readSize += benchReadFile(ss2.str());
                }

                readTime = getTime() - readTime;
                cout << setw(12) << fileSizeToString(readSize) << "," << setw(12) << right << fileSpeedToString(readSize, readTime);
                cout.flush();

                // Perform remove test
                deleteTime = getTime();
                for (U_LONG i = 0; i < fCount; ++i)
                {
                    stringstream ss1, ss2;
                    ss1 << DIR << fName1 << i;
                    ss2 << DIR << fName2 << i;
                    fileCleanerA.unregisterFile(ss1.str());
                    fileCleanerB.unregisterFile(ss2.str());
                    remove(ss1.str().c_str());
                    remove(ss2.str().c_str());
                }
                deleteTime      = getTime() - deleteTime;
                completeTime    = getTime() - completeTime;
                completeSeconds = time(0)   - completeSeconds;

                cout << " | deleted in:" << setw(7) << deleteTime << " ms";
                cout << " | completed in:" << setw(7) << completeTime << " ms,"<<completeSeconds<< " s";
                cout << endl;
            }
        }

    } catch (string s)
    {
        cout << "Test failed, some files may still remain on the disk !"
                << endl;
        cout << "Exception: " << s << endl;
        return 0;
    }

    return 0;

}
